915F - Imbalance Value of a Tree - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;

long long n, mx, mn, u, v, w[1000007];
pair<long long, long long> a[1000007];
vector<long long> adj[1000007];

struct DSU{
    vector<int> par, sz;

    void make_set(int n){
        par.resize(n + 1);
        sz.resize(n + 1);
        for(int i = 1; i <= n; i++){
            par[i] = i;
            sz[i] = 1;
        }
    }

    int find_set(int v){
        if(par[v] == v) return v;
        return par[v] = find_set(par[v]);
    }

    bool union_sets(int u, int v){
        u = find_set(u);
        v = find_set(v);
        if(u == v) return false;
        if(sz[u] < sz[v]) swap(u, v);
        par[v] = u;
        sz[u] += sz[v];
        return true;
    }
} dsu;

void solveMax(){
    dsu.make_set(n);
    sort(a + 1, a + n + 1);
    for(int i = 1; i <= n; i++){
        int u = a[i].second;
        long long last = 0, szu = dsu.sz[dsu.find_set(u)];
        for(int j = 0; j < (int) adj[u].size(); j++){
            int v = adj[u][j];
            if(w[v] <= w[u]){
                long long szv = dsu.sz[dsu.find_set(v)];
                if(dsu.union_sets(u, v)){
                    mx += szu * szv * a[i].first;
                    mx += last * szv * a[i].first;
                    last += szv;
                }
            }
        }
    }
}

bool cmp(const pair<long long, long long> &a, const pair<long long, long long> &b){
    return a.first > b.first;
}

void solveMin(){
    dsu.make_set(n);
    sort(a + 1, a + n + 1, cmp);
    for(int i = 1; i <= n; i++){
        int u = a[i].second;
        long long last = 0, szu = dsu.sz[dsu.find_set(u)];
        for(int j = 0; j < (int) adj[u].size(); j++){
            int v = adj[u][j];
            if(w[v] >= w[u]){
                long long szv = dsu.sz[dsu.find_set(v)];
                if(dsu.union_sets(u, v)){
                    mn += szu * szv * a[i].first;
                    mn += last * szv * a[i].first;
                    last += szv;
                }
            }
        }
    }
}

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(0); cout.tie(0);
    //freopen("marisa.inp", "r", stdin);
    //freopen("marisa.out", "w", stdout);
    cin >> n;
    for(int i = 1; i <= n; i++){
        cin >> a[i].first;
        a[i].second = i;
        w[i] = a[i].first;
    }
    for(int i = 1; i < n; i++){
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    solveMax();
    solveMin();

    cout << mx - mn << endl;
}